#include "ReconstructByKinectCommon.h"
#include <iostream>
#include <sstream>
#include <fstream>

using namespace std;

const int PATT_ROWS = 5;
const int PATT_COLUMNS = 7;
const int PATT_CORNERS = PATT_ROWS * PATT_COLUMNS;
const CvSize2D32f SQUARE_SIZE = cvSize2D32f (10.0f, 10.0f);
const CvSize FRAME_SIZE = cvSize (640, 480);
const int MAX_GROUPS = 20;

IplImage *g_pCurrentFrame = NULL;
IplImage *g_pattResultImage = NULL;
CvPoint3D32f g_patternCorners[PATT_ROWS][PATT_COLUMNS];
CvPoint2D32f g_imgCornersGroups[PATT_CORNERS * MAX_GROUPS];
string g_imageIndexList[MAX_GROUPS];
int g_groups = 0;
string g_testName;
string g_rootDir;
CvFont g_font;
ofstream g_fout;

void showRecognizeResult (IplImage *image, CvPoint2D32f *corners)
{
	cvCopy (image, g_pattResultImage);
	cvDrawChessboardCorners (g_pattResultImage, cvSize (PATT_COLUMNS, PATT_ROWS),
			corners, PATT_CORNERS, 1);
	ostringstream sout;
	sout << "Group: " << g_imageIndexList[g_groups - 1];
	cvPutText (g_pattResultImage, sout.str ().c_str (), cvPoint (5, 20),
			&g_font, cvScalar (0, 0, 255));
}

bool recognizePattern (IplImage *image, CvPoint2D32f *corners)
{
	if (!image)
		return false;
	int cornerCount;
	int result = cvFindChessboardCorners (image,
			cvSize (PATT_COLUMNS, PATT_ROWS), corners, &cornerCount);
	if (result == 1) {
// 		cvFindCornerSubPix (image, corners, cornerCount,
// 				cvSize (5, 5), cvSize (1, 1),
// 				cvTermCriteria (CV_TERMCRIT_ITER | CV_TERMCRIT_EPS, 10, 0.02));
		if (glm::length (glm::vec2 (corners[0].x, corners[0].y)) >
				glm::length (glm::vec2 (corners[PATT_ROWS * PATT_COLUMNS - 1].x,
				corners[PATT_ROWS * PATT_COLUMNS - 1].y))) {
			for (int i = 0; i < PATT_ROWS * PATT_COLUMNS / 2; ++i) {
				CvPoint2D32f temp = corners[i];
				corners[i] = corners[PATT_ROWS * PATT_COLUMNS - 1 - i];
				corners[PATT_ROWS * PATT_COLUMNS - 1 - i] = temp;
			}
		}
	}
	return result == 1;
}

void showCalibrationResult (CvMat *pIntrinsicMat, CvMat *pDistCoeffsMat)
{
	cout << "*******************Calibration Result********************" << endl;
	cout << "Intrinsic Matrix: \n";
	for(int v = 0; v < 3; ++v) {
		for (int u = 0; u < 3; ++u) {
			float Val = CV_MAT_ELEM (*pIntrinsicMat, float, v, u);
			cout << Val << "\t";
		}
		cout << endl;
	}
	cout << "Dist Ceoffs: ";
	for(int i = 0; i < 5; ++i) {
		float Val = CV_MAT_ELEM (*pDistCoeffsMat, float, i, 0);
		cout << Val;
		if (i != 4)
			cout << ", ";
		else
			cout << endl;
	}
}

void saveReprojectResult (int group, CvMat *corners)
{
	string imageFile = g_rootDir + "Input_" + g_imageIndexList[group] + ".bmp";
	IplImage *pOrigImage = cvLoadImage (imageFile.c_str ());
	if (pOrigImage == NULL)
		return;
	cvDrawChessboardCorners (pOrigImage, cvSize (PATT_COLUMNS, PATT_ROWS),
			&CV_MAT_ELEM (*corners, CvPoint2D32f, 0, 0), PATT_CORNERS, 1);
	ostringstream sout;
	sout << "Group: " << g_imageIndexList[group];
	cvPutText (pOrigImage, sout.str ().c_str (), cvPoint (5, 20),
			&g_font, cvScalar (255, 0, 0));
	
	imageFile = g_rootDir + "ReprojectResult_" +
			g_imageIndexList[group] + ".bmp";
	cvSaveImage (imageFile.c_str (), pOrigImage);
	cvReleaseImage (&pOrigImage);
}

void saveUndisortResult (int group, const CvMat* pCameraMat,
		const CvMat* pDistCoeffsVecM)
{
	string imageFile = g_rootDir + "Input_" + g_imageIndexList[group] + ".bmp";
	IplImage *pOrigImage = cvLoadImage (imageFile.c_str ());
	IplImage *pUndistortedImage = cvCreateImage (cvGetSize (pOrigImage),
			pOrigImage->depth, pOrigImage->nChannels);
	if (pOrigImage == NULL)
		return;
	cvUndistort2 (pOrigImage, pUndistortedImage, pCameraMat, pDistCoeffsVecM);
	imageFile = g_rootDir + "UndistortedResult_" +
		g_imageIndexList[group] + ".bmp";
	cvSaveImage (imageFile.c_str (), pUndistortedImage);
	cvReleaseImage (&pOrigImage);
	cvReleaseImage (&pUndistortedImage);
}

void verifyCalibrationResult (
		CvMat *pObjectCornersListMat, CvMat *pImageCornersListMat,
		CvMat *pIntrinsicMat, CvMat *pDistCoeffsMat,
		CvMat *pRVecsMat, CvMat *pTVecsMat)
{
	double totalErr = 0.0, err = 0.0;
	CvMat objCornersM;
	CvMat imgCornersM;
	CvMat rVec;
	CvMat tVec;

	CvMat *pProjectedObjCornersM = cvCreateMat (PATT_CORNERS, 1, CV_32FC2);

	cout << "********************Result Error*********************" << endl;
	cout << "Error (per image): " << endl;

	for (int i = 0; i < g_groups; i++) {
		cvGetRows (pObjectCornersListMat, &objCornersM,
				i * PATT_CORNERS, (i + 1) * PATT_CORNERS);
		cvGetRows (pImageCornersListMat, &imgCornersM,
				i * PATT_CORNERS, (i + 1) * PATT_CORNERS);
		cvGetRow (pRVecsMat, &rVec, i);
		cvGetRow (pTVecsMat, &tVec, i);

		cvProjectPoints2 (&objCornersM, &rVec, &tVec,
				pIntrinsicMat, pDistCoeffsMat, pProjectedObjCornersM);
		saveReprojectResult (i, pProjectedObjCornersM);
		saveUndisortResult (i, pIntrinsicMat, pDistCoeffsMat);
		err = cvNorm (&imgCornersM, pProjectedObjCornersM, CV_L1);
		totalErr += err /= PATT_ROWS * PATT_COLUMNS;
		cout << i << ": " << err << " pixels" << endl;
	}
	cout << "Average error: " << totalErr / g_groups << " pixels" << endl;
}


void calculateAndVerifyIntrinsicParam ()
{
	cout << "Start Calibration...";
	if (g_groups <= 5) {
		cout << "Failed: No enough data!" << endl;
		return;
	}
	int totalCorners = g_groups * PATT_CORNERS;
	
	// Input Mats
	CvMat *pImageCornersGroupsM = cvCreateMat (totalCorners, 1, CV_32FC2);
	for (int i = 0; i < totalCorners; ++i)
		CV_MAT_ELEM (*pImageCornersGroupsM, CvPoint2D32f, i, 0) = g_imgCornersGroups[i];

	CvMat *pObjCornersGroupsM = cvCreateMat (totalCorners, 1, CV_32FC3);
	for (int g = 0, i = 0; g < g_groups; ++g)
		for (int v = 0; v < PATT_ROWS; ++v)
			for (int u = 0; u < PATT_COLUMNS; ++u, ++i)
				CV_MAT_ELEM (*pObjCornersGroupsM, CvPoint3D32f, i, 0) =
						g_patternCorners[v][u];

	CvMat *pCornersCountVecM = cvCreateMat (g_groups, 1, CV_32SC1);
	for (int i = 0; i < g_groups; ++i)
		CV_MAT_ELEM (*pCornersCountVecM, int, i, 0) = PATT_CORNERS;
	
	CvMat *pInitIntrinsicMat = cvCreateMat (3, 3, CV_32FC1);

	// Result Mats
	CvMat *pIntrinsicMat = cvCreateMat (3, 3, CV_32FC1);
	CvMat *pDistCoeffsMat = cvCreateMat (5, 1, CV_32FC1);
	CvMat *pRVecsM = cvCreateMat (g_groups, 1, CV_32FC3);
	CvMat *pTVecsM = cvCreateMat (g_groups, 1, CV_32FC3);

	cvCalibrateCamera2 (pObjCornersGroupsM, pImageCornersGroupsM,
			pCornersCountVecM, FRAME_SIZE, pIntrinsicMat, pDistCoeffsMat,
			pRVecsM, pTVecsM);
	cout << "Completed!" << endl;

	showCalibrationResult (pIntrinsicMat, pDistCoeffsMat);
	verifyCalibrationResult (pObjCornersGroupsM, pImageCornersGroupsM, 
			pIntrinsicMat, pDistCoeffsMat, pRVecsM, pTVecsM);
}


void createPatternCorners ()
{
	for (int v = 0; v < PATT_ROWS; ++v)
		for (int u = 0; u < PATT_COLUMNS; ++u)
			g_patternCorners[v][u] = cvPoint3D32f (u * SQUARE_SIZE.width,
					v * SQUARE_SIZE.height, 0);
}

void saveRecognitionResult ()
{
	string imageFile = g_rootDir + "RecognitionResult_" +
			g_imageIndexList[g_groups - 1] + ".bmp";
	cvSaveImage (imageFile.c_str (), g_pattResultImage);
}

void saveCurrentFrame ()
{
	string imageFile = g_rootDir + "Input_" +
			g_imageIndexList[g_groups - 1] + ".bmp";
	cvSaveImage (imageFile.c_str (), g_pCurrentFrame);
	g_fout << imageFile << endl;
}

bool useCurrentFrame (string imageIndex)
{
	bool isSuccess = recognizePattern (g_pCurrentFrame,
			&g_imgCornersGroups[g_groups * PATT_CORNERS]);
	if (isSuccess) {
		g_imageIndexList[g_groups] = imageIndex;
		++g_groups;
		showRecognizeResult (g_pCurrentFrame,
				&g_imgCornersGroups[(g_groups - 1) * PATT_CORNERS]);
		saveRecognitionResult ();
		return true;
	}
	return false;
}

void cleanAndExit ()
{
	system ("pause");
	exit (0);
}

void processFromSensor (SensorDataInput *pSensor)
{
	int key = -1;
	g_fout.open ((g_rootDir + "list.txt").c_str ());
	g_pCurrentFrame = pSensor->getRGBImage ();
	while (true) {
		if (key == 27 || g_groups == MAX_GROUPS || !pSensor->getNewFrame ()) {
			calculateAndVerifyIntrinsicParam ();
			cleanAndExit ();
		}
		cvShowImage ("Sensor", g_pCurrentFrame);
		key = cvWaitKey (30);
		if (key == ' ') {
			ostringstream sout;
			sout << g_groups - 1;
			if (useCurrentFrame (sout.str ())) {
				cvShowImage ("Pattern Corners", g_pattResultImage);
				saveCurrentFrame ();
			}
		}
	}
}

void processFromSeq (ifstream &fIn)
{
	while (true) {
		string imageIndex;
		getline (fIn, imageIndex);
		if (!fIn || imageIndex == "")
			break;
		string imageFile = g_rootDir + "Input_" + imageIndex + ".bmp";
		g_pCurrentFrame = cvLoadImage (imageFile.c_str ());
		if (g_pCurrentFrame == NULL) {
			cerr << "Load image fail: " << imageFile << endl;
			throw exception ("Load image fail");
		}
		cout << "Processing: " << imageFile << endl;
		useCurrentFrame (imageIndex);
		cvReleaseImage (&g_pCurrentFrame);
	}
	calculateAndVerifyIntrinsicParam ();
	cleanAndExit ();
}

#define FROM_KINECT 0
#define FROM_RECORD 1
#define FROM_SEQ 2
int main (int argc, char* argv[]) {
	if (argc < 3) {
		printf ("Command line argument error.\n");
		return -1;
	}
	g_testName = argv[1];
	g_rootDir = "..\\Data\\Calibration\\";
	g_rootDir += g_testName + "\\";
	int dataSource;
	if (strcmp (argv[2], "kinect") == 0 && argc >= 3)
		dataSource = FROM_KINECT;
	else if (strcmp (argv[2], "file") == 0 && argc >= 3)
		dataSource = FROM_RECORD;
	else if (strcmp (argv[2], "seq") == 0)
		dataSource = FROM_SEQ;
	else {
		printf ("Sensor must be kinect or file.\n");
		return -1;
	}
	try {
		cvInitFont (&g_font, CV_FONT_HERSHEY_PLAIN, 1.0f, 1.0f);
		createPatternCorners ();
		g_pattResultImage = cvCreateImage (cvSize (640, 480), 8, 3);
		switch (dataSource) {
		case FROM_KINECT: 
			processFromSensor (
					new SensorDataKinectInput (string ("Sensor")));
			break;
		case FROM_RECORD:
			processFromSensor (
					new SensorDataFileInput (string ("Sensor"), argv[3]));
			break;
		case FROM_SEQ:
			processFromSeq (ifstream ((g_rootDir + "list.txt").c_str ()));
			break;
		}
	} catch (...) {
		printf ("Exception caught\n");
		system ("pause");
	}
	return 0;

}